# qwenvl3_search_gpu_parallel.py
from demo_qwen3_api import get_qwen3_client, get_ans_dict_by_search, get_non_empty_keys, get_matched_info_by_jsondict

qwen_client = None
df = None
empty_emb = None
template_json_str = None

import concurrent.futures
import numpy as np
import heapq
import json
import torch
import torch.nn.functional as F

import asyncio
import httpx

async def check_satisfy_sim_api_async(
    target_json_dict_str,
    non_empty_keys,
    search_joined_values_list,
    search_embs,
    threshold=0.3,
    device="cpu",
    api_url="http://localhost:20022/satisfy_sim"
):
    """
    异步调用 Flask 检查 API。
    参数:
        target_json_dict_str: str, 待判断的 json 字符串
        non_empty_keys: list                            ['basic_info', 'season']
        search_joined_values_list: list                 ['spring']
        search_embs: list (向量列表，通常是嵌入)
        threshold: float
        device: str
        api_url: str, API 地址
    返回:
        satisfy: bool
        sim_list: list
    """
    search_embs = [emb.tolist() if hasattr(emb, 'tolist') else emb 
                                for emb in search_embs]
    payload = {
        "target_json_dict_str": target_json_dict_str,
        "non_empty_keys": non_empty_keys,
        "search_joined_values_list": search_joined_values_list,
        "search_embs": search_embs,
        "threshold": threshold,
        "device": device
    }
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            resp = await client.post(api_url, json=payload)
            resp.raise_for_status()
            result = resp.json()
            return result.get("satisfy"), result.get("sim_list")
    except Exception as e:
        print(f"Async API call error: {e}")
        return False, []


async def search_and_show(search_txt, threshold=0.5):
    import pandas as pd
    global qwen_client, df,empty_emb,template_json_str
    # 读取csv
    if df is None:
        csv_file_path = '/mnt/nas/shengjie/qdrant_data/qwenvl3_data/qwenvl3_cloth_data2600.csv'
        # csv_file_path = '/mnt/nas/shengjie/qdrant_data/qwenvl3_data/qwenvl3_cloth_data_api.csv'
        df = pd.read_csv(csv_file_path)
    if template_json_str is None:
        # 读取模板json字符串
        template_json_path = 'demo_qwenvl_clothing.json'
        with open(template_json_path, 'r', encoding='utf-8') as f:
            template_json_str = f.read()
    # 获取qwen3客户端
    if qwen_client is None:
        qwen_client = get_qwen3_client()
    # 获取推理结果
    ans = get_ans_dict_by_search(qwen_client, template_json_str, search_txt)
    import json
    try:
        ans_dict = json.loads(ans)
    except Exception as e:
        print("Failed to parse ans as JSON:", e)
        ans_dict = None
    # 获得非空key链
    non_empty_keys = get_non_empty_keys(ans_dict)

    # ----------- (1) 筛选 key，仅保留指定key链 ------------

    # 定义允许的key链前缀(只判断到某一级即可)
    allowed_prefixes = [
        ['basic_info', 'season'],
        ['basic_info', 'gender'],
        ['design', 'color', 'primary'],
        ['design', 'color', 'secondary'],
        ['design', 'color', 'pattern'],
        ['design', 'category', 'main'],
        ['design', 'category', 'sub'],
        ['design', 'sleeve', 'type'],
        ['design', 'sleeve', 'cuff'],
        ['design', 'sleeve', 'count'],
        ['design', 'collar', 'type'],
        ['design', 'pockets'],         # pockets是整个list,后续get_value_by_keys能拿到
        ['material', 'composition'],
        ['decoration'],                # 只要一级即可，后续靠get_value_by_keys拉取全部
        ['visual'],
    ]

    # 给定实际key链，判定是否属于允许集合
    def is_key_allowed(keys):
        for prefix in allowed_prefixes:
            if len(keys) < len(prefix):
                continue
            # 只要前缀对上即可
            if keys[:len(prefix)] == prefix:
                return True
        return False

    print('[before] non_empty_keys:',non_empty_keys)
    # 对non_empty_keys做一次筛选
    non_empty_keys = [keys for keys in non_empty_keys if is_key_allowed(keys)]
    print('[sifted] non_empty_keys:',non_empty_keys)
    

    # get_matched_info_by_jsondict 这个函数应该全局可用
    from util_for_clipvitb32 import clip_vit_b32, get_embedding
    import torch
    import torch.nn.functional as F
    import heapq
    clip_model = clip_vit_b32()
    if empty_emb is None:
        empty_txt = ''
        empty_emb = get_embedding(clip_model, empty_txt)    
    search_joined_values_list = get_matched_info_by_jsondict(ans_dict, non_empty_keys)
    search_joined_values_embs = [get_embedding(clip_model, item) 
                                    for item in search_joined_values_list]
    search_joined_values = '&&'.join(search_joined_values_list)                                    

    topk = 10
    heap = []

    # 准备所有API调用任务
    tasks = []
    for idx, row in df.iterrows():
        target_json_dict_str = row.get('desc_json', None)
        img_path = row.get('img_path', '')
        
        tasks.append(
            check_satisfy_sim_api_async(
                target_json_dict_str,
                non_empty_keys,
                search_joined_values_list,
                search_joined_values_embs,
                threshold
            )
        )
        print(f'\r idx: {idx}', end='')
    print('\n**API任务全部创建完成**')
    
    # INSERT_YOUR_CODE
    # 尝试使用 asyncio.as_completed 收集所有结果（替代 gather 的用法）
    # results = []
    # import time
    # # 记录每一次迭代用时
    # record_time = []
    # for fut in asyncio.as_completed(tasks):
    #     time_start = time.time()
    #     try:
    #         result = await fut
    #     except Exception as e:
    #         result = e
    #     time_cost = time.time() - time_start
    #     record_time.append(time_cost)
    #     results.append(result)
    #     print(f'\r id : {len(results)}',end='')
    # if record_time:
    #     max_time = max(record_time)
    #     if len(record_time) > 1:
    #         avg_time = sum(record_time) - max_time
    #         avg_time = avg_time / (len(record_time)-1)
    #     else:
    #         avg_time = 0.0
    #     print(f"\n最大用时: {max_time:.5f} s, 其余平均用时: {avg_time:.5f} s")
    # print('\n**执行API任务 (as_completed收集)**')

    # 并发执行所有API调用(等所有的请求处理结束)
    results = await asyncio.gather(*tasks, return_exceptions=True)
    print('**执行API任务**')
    

    # 处理结果
    for idx, (result, row) in enumerate(zip(results, df.iterrows())):
        _, row_data = row
        img_path = row_data.get('img_path', '')
        
        # 处理可能的异常
        if isinstance(result, Exception):
            print(f"API call failed for index {idx}: {result}")
            continue
        
        satisfy, sim_list = result
        if satisfy:
            sum_sim = sum(sim_list)
            joined_values_list = get_matched_info_by_jsondict(
                json.loads(row_data['desc_json']), 
                non_empty_keys
            )
            item = (
                sum_sim,
                [f'{sim:.4f}' for sim in sim_list],
                idx,
                img_path,
                '&&'.join(joined_values_list)
            )
            if len(heap) < topk:
                heapq.heappush(heap, item)
            else:
                if sum_sim > heap[0][0]:
                    heapq.heappushpop(heap, item)

        print(f'\r idx: {idx}', end='')
    print('\n结果处理完成')

    # 排序，返回info和图片路径
    heap_sorted = sorted(heap, key=lambda x: -x[0])
    res_list = []
    img_list = []
    from utils.util_flux import process_img_1024
    for sim, sim_list, idx, path, tar_value_str in heap_sorted:
        res_list.append( f"TOPK sum_sim={sim:.4f} search_sim={sim_list}  \
                         idx={idx} path= {path} desc={tar_value_str}")
        img_list.append(
            process_img_1024(path)
        )

    # 返回search_txt, 图片列表，内容列表
    # non_empty_keys search_joined_values_list
    search_txt_list = [f'{k}:{v}' for k,v in zip(non_empty_keys, search_joined_values_list)]
    search_txt = '\n'.join(search_txt_list)
    return search_txt, img_list, '\n'.join(res_list)

async def async_search_wrapper(search_txt: str, threshold: float = 0.5):
    """异步包装器，供Gradio直接调用"""
    try:
        return await search_and_show(search_txt, threshold)
    except Exception as e:
        print(f"Search error: {e}")
        return "Error", [], "Search failed"


def start_gradio(port):
    import gradio as gr
    with gr.Blocks() as demo:
        gr.Markdown("# QwenVL3 Gradio Demo (以文搜图)")
        with gr.Row():
            # with gr.Column():
            search_box = gr.Textbox(label="输入搜索描述", placeholder="如：春秋两季黑色大衣，两个口袋", lines=2)
        with gr.Row():
            threshold_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.01, label="阈值 threshold", interactive=True)
            submit_btn = gr.Button("开始搜索")
            # with gr.Column():
        with gr.Row():
            output_search = gr.Textbox(label="搜索内容", interactive=False)
        with gr.Row():
            gallery = gr.Gallery(label="TopK图片预览", columns=5, height="auto")
        with gr.Row():
            output_list = gr.Textbox(label="TopK内容", lines=10, interactive=False)
        submit_btn.click(
            async_search_wrapper,
            inputs=[search_box, threshold_slider],
            outputs=[output_search, gallery, output_list]
        )
    demo.launch(server_name="0.0.0.0", server_port=port)

# 若本py被直接调用则自动启动界面
if __name__ == '__main__':
    import os
    os.environ['GRADIO_TEMP_DIR'] = '/mnt/nas/shengjie/tmp'
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--cuda', type=str, default='2', help='CUDA device id')
    parser.add_argument('-p', '--port', type=int, default=20023, help='CUDA device id')
    args, unknown = parser.parse_known_args()
    os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda

    start_gradio(args.port)